"""
主入口文件，用于训练和运行LSTM模型
"""
import argparse
import os
import sys

# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from clip_model.train_clip import train_clip_model
from clip_model.train_lstm import train_lstm_model, train_clip_lstm_model


def main():
    parser = argparse.ArgumentParser(description="Train and run LSTM models")
    parser.add_argument("--model", type=str, default="clip_lstm", 
                        choices=["clip", "lstm", "clip_lstm"],
                        help="Model to train: clip, lstm, or clip_lstm")
    parser.add_argument("--epochs", type=int, default=3, help="Number of epochs")
    parser.add_argument("--batch_size", type=int, default=16, help="Batch size")
    parser.add_argument("--lr", type=float, default=5e-4, help="Learning rate")
    
    args = parser.parse_args()
    
    print(f"Training {args.model} model...")
    print(f"Epochs: {args.epochs}, Batch size: {args.batch_size}, Learning rate: {args.lr}")
    
    if args.model == "clip":
        # 训练CLIP模型
        train_clip_model(
            epochs=args.epochs,
            batch_size=args.batch_size,
            learning_rate=args.lr,
            save_path="clip_model.pth"
        )
    elif args.model == "lstm":
        # 训练LSTM模型
        train_lstm_model(
            epochs=args.epochs,
            batch_size=args.batch_size,
            learning_rate=args.lr,
            save_path="lstm_model.pth"
        )
    elif args.model == "clip_lstm":
        # 训练CLIP-LSTM混合模型
        # 首先确保有CLIP模型
        if not os.path.exists("clip_model.pth"):
            print("CLIP model not found. Training CLIP model first...")
            train_clip_model(
                epochs=args.epochs,
                batch_size=args.batch_size,
                learning_rate=args.lr,
                save_path="clip_model.pth"
            )
        
        # 训练CLIP-LSTM混合模型
        train_clip_lstm_model(
            clip_model_path="clip_model.pth",
            epochs=args.epochs,
            batch_size=args.batch_size//2,  # 使用较小的批次大小
            learning_rate=args.lr/10,  # 使用较小的学习率
            save_path="clip_lstm_model.pth"
        )
    
    print(f"Training {args.model} model completed!")


if __name__ == "__main__":
    main()